from torch.utils.data import Dataset, DataLoader
import numpy as np
import os
from glob import glob
import torch
from training.utils.common_util import set_random_seed, split_files_by_kfold, split_files_by_subject_kfold
from training.utils.stats_util import StatsUtil, get_file_stages_from_stats
from training.utils.data_util import (
    LRUCache,
    process_sample,
    collate_dict,
    resolve_path,
    optimize_dataloader_worker_config,
    create_optimized_dataloader
)


class SleepDataset(Dataset):
    def __init__(self, files, target_sample_rate=100, do_normalize=True, stats_file=None, cache_size=10):
        self.files = files
        self.target_sample_rate = target_sample_rate
        self.do_normalize = do_normalize
        self.stats_file = stats_file
        self.cache_size = cache_size  # 文件缓存大小限制
        self.stages_counts = {}  # 存储每个文件的阶段数量
        self.normalization_stats = {}
        
        # 初始化阶段数量统计
        self._initialize_stages_counts()
        
        # 创建全局索引映射：将全局索引映射到(文件索引, 阶段索引)对
        self.global_to_file_mapping = []
        self._create_global_index_mapping()
        
        # 使用工具类中的LRUCache实现
        self.file_cache = LRUCache(capacity=cache_size)
        
        # 如果需要标准化，加载标准化统计信息
        if self.do_normalize:
            self._load_normalization_stats()
    
    def _create_global_index_mapping(self):
        """
        创建全局索引到(文件索引, 阶段索引)的映射
        用于快速定位特定阶段所属的文件
        使用已经通过get_file_stages_from_stats获取的阶段数量，如果失败则直接加载文件
        """
        for file_idx, file_path in enumerate(self.files):
            # 首先尝试使用已在_initialize_stages_counts中获取的阶段数量
            stages_count = self.stages_counts.get(file_path, 0)
            
            # 如果阶段数量为0，尝试直接加载文件获取阶段数量
            if stages_count == 0:
                try:
                    print(f"正在直接处理文件: {file_path}")
                    # 直接加载文件以确定阶段数量
                    data = np.load(file_path)
                    
                    # 处理npz文件格式
                    if isinstance(data, np.lib.npyio.NpzFile):
                        # 打印文件中包含的键，帮助调试
                        print(f"文件 {file_path} 包含的键: {list(data.keys())}")
                        
                        # 检查不同的可能数据格式
                        if 'y' in data:
                            stages_count = len(data['y'])
                        elif 'x' in data:
                            stages_count = len(data['x'])
                        else:
                            # 尝试找到第一个合适的数组
                            for key in data.keys():
                                arr = data[key]
                                if isinstance(arr, np.ndarray) and len(arr.shape) > 0:
                                    stages_count = len(arr)
                                    print(f"使用键 '{key}' 确定阶段数量: {stages_count}")
                                    break
                        # 关闭npz文件
                        data.close()
                    
                    # 更新阶段数量统计
                    self.stages_counts[file_path] = stages_count
                    
                except Exception as e:
                    print(f"加载文件时出错 {file_path}: {str(e)}")
                    stages_count = 0
            
            # 为该文件的每个阶段创建映射
            for stage_idx in range(stages_count):
                self.global_to_file_mapping.append((file_idx, stage_idx))
        
        # 记录映射创建情况
        print(f"成功创建全局索引映射，总阶段数: {len(self.global_to_file_mapping)}")
        
        # 验证数据集是否有有效样本
        if len(self.global_to_file_mapping) == 0:
            print("警告: 没有找到有效的数据样本，数据集为空")
    
    def _load_normalization_stats(self):
        """
        从统计文件加载标准化所需的均值和标准差，不直接操作数据集
        """
        self.normalization_stats = StatsUtil.get_normalization_stats_from_stats(self.stats_file)
        if not self.normalization_stats:
            print("无法从统计文件加载标准化统计信息，使用默认值")
            self.normalization_stats = {'mean': [0.0], 'std': [1.0]}
    
    def _load_precomputed_stats(self, file_path):
        """
        从统计文件加载预先计算好的统计信息
        """
        return StatsUtil.get_file_stats(file_path, self.stats_file)
    
    def _initialize_stages_counts(self):
        """
        初始化阶段数量统计，使用stats_util中的辅助函数高效获取所有文件的阶段数量
        确保正确统计每个文件的阶段数量，为全局索引映射做准备
        """
        # 使用stats_util中的get_file_stages_from_stats函数批量获取阶段数量
        # 该函数提供了更好的错误处理、缓存机制和回退计算功能
        stages_counts_list = get_file_stages_from_stats(self.stats_file, self.files, fallback_to_compute=True)
        
        # 构建stages_counts字典
        for file_path, stages_count in zip(self.files, stages_counts_list):
            self.stages_counts[file_path] = stages_count
        
        # 记录统计信息加载情况
        valid_counts = sum(1 for count in self.stages_counts.values() if count > 0)
        print(f"成功加载/计算 {valid_counts}/{len(self.files)} 个文件的阶段数量")
        if self.stats_file:
            print(f"使用统计文件: {os.path.basename(self.stats_file)}")
    
    def _process_sample(self, data):
        """
        处理单个样本
        使用工具类中的process_sample函数
        """
        mean = np.array(self.normalization_stats['mean']) if self.do_normalize and self.normalization_stats else None
        std = np.array(self.normalization_stats['std']) if self.do_normalize and self.normalization_stats else None
        return process_sample(data, mean=mean, std=std, add_channel_dim=True)
    
    def __len__(self):
        """
        返回数据集的大小 - 所有文件中的睡眠阶段总数
        """
        return len(self.global_to_file_mapping)
    
    def __getitem__(self, idx):
        """
        获取指定索引的数据样本 - 单个睡眠阶段
        使用工具类中的LRUCache提高加载效率
        """
        # 根据全局索引找到对应的文件和阶段
        file_idx, stage_idx = self.global_to_file_mapping[idx]
        file_path = self.files[file_idx]
        
        # 加载数据 - 使用LRUCache工具类
        try:
            # 检查文件是否在缓存中
            data = self.file_cache.get(file_path)
            
            # 未命中缓存，加载文件并添加到缓存
            if data is None:
                data = np.load(file_path)
                self.file_cache.put(file_path, data)
            
            # 获取单个阶段的数据
            x = data['x'][stage_idx]  # 获取单个阶段的特征数据
            y = data['y'][stage_idx]  # 获取单个阶段的标签数据
            
            # 进行必要的预处理
            x_processed = self._process_sample(x)
            
            # 转换为PyTorch张量
            x_tensor = torch.from_numpy(x_processed).float()
            y_tensor = torch.tensor(y).long()
            
            return {
                'data': x_tensor,
                'label': y_tensor,
                'filename': os.path.basename(file_path),
                'stage_idx': stage_idx
            }
        except Exception as e:
            print(f"加载文件 {file_path} 的阶段 {stage_idx} 失败: {str(e)}")
            # 返回默认值，避免程序崩溃
            return {'data': torch.zeros((3000,), dtype=torch.float32), 
                    'label': torch.zeros(1, dtype=torch.long),
                    'filename': os.path.basename(file_path),
                    'stage_idx': stage_idx}

def create_dataloaders(train_files, val_files, batch_size, num_workers=None, pin_memory=True, 
                      target_sample_rate=100, do_normalize=True, stats_file=None, train_stats_only=False):
   
    # 使用工具类中的函数计算最优worker数量和缓存大小
    total_files = len(train_files) + (len(val_files) if val_files else 0)
    num_workers, cache_size = optimize_dataloader_worker_config(total_files, default_workers=num_workers)
    print(f"优化配置: {num_workers}个worker, 缓存大小{cache_size}")
    
    # 检查数据集是否为空
    if not train_files:
        print("警告: 训练集文件列表为空，无法创建有效的数据加载器")
        return None, None

    # 加载训练集的标准化统计信息（如果需要）
    train_normalization_stats = None
    if do_normalize and stats_file and train_stats_only:
        # 加载训练集的统计信息
        train_normalization_stats = StatsUtil.get_normalization_stats_from_stats(stats_file)
        print(f"使用训练集统计信息进行标准化: {train_normalization_stats}")
    else:
        train_normalization_stats = None
        print("跨域训练模式：只使用训练集的均值和标准差进行标准化")
    
    # 训练集和验证集分别标准化，避免数据泄露，同时启用文件缓存
    train_dataset = SleepDataset(train_files, target_sample_rate=target_sample_rate, 
                               do_normalize=do_normalize, stats_file=stats_file, cache_size=cache_size)
    
    # 只在val_files非空时创建验证集
    val_dataset = None
    if val_files:
        val_dataset = SleepDataset(val_files, target_sample_rate=target_sample_rate, 
                                 do_normalize=do_normalize, stats_file=stats_file, cache_size=cache_size)
        # 如果启用了train_stats_only，为验证集设置与训练集相同的标准化统计信息
        if train_stats_only and train_normalization_stats:
            val_dataset.normalization_stats = train_normalization_stats
    
    # 使用工具类中的collate_dict函数，无需重复定义
    
    # 训练集数据加载器 - 使用工具类中的优化函数
    train_loader = create_optimized_dataloader(
        train_dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=num_workers,
        pin_memory=pin_memory,
        drop_last=True,
        collate_fn=collate_dict
    )
    
    val_loader = None
    if val_dataset:
        # 验证集数据加载器 - 使用优化函数，worker数量减半
        val_loader = create_optimized_dataloader(
            val_dataset,
            batch_size=batch_size,
            shuffle=False,
            num_workers=max(1, num_workers // 2),  # 验证集使用较少的worker
            pin_memory=pin_memory,
            drop_last=False,
            collate_fn=collate_dict
        )
    
    print(f"创建数据加载器: batch_size={batch_size}, 训练集worker={num_workers}, 验证集worker={max(1, num_workers // 2) if val_loader else 0}")
    return train_loader, val_loader


# 新增 max_files 参数：限制总文件数；add_random 参数：是否随机选择文件（默认False取前N个）
def get_kfold_files(data_dir, n_folds, max_files=None, add_random=False):
    """
    获取K折交叉验证的数据文件列表
    """
    # 设置随机种子，确保文件选择的可复现性
    set_random_seed(42)
    
    # 使用工具类中的resolve_path函数处理路径
    data_dir = resolve_path(data_dir)

    print(f"Using data directory: {data_dir}")  # 调试信息
    all_files = glob(os.path.join(data_dir, "*.npz"))
    print(f"Found {len(all_files)} files")  # 调试信息
    all_files.sort()  # 确保文件顺序固定

    # 限制文件数量
    if max_files is not None and max_files > 0:
        if add_random:
            # 随机选择 max_files 个文件（避免总是用前N个导致的数据偏差）
            import random  # 在这里局部导入以避免潜在的命名冲突
            all_files = random.sample(all_files, min(max_files, len(all_files)))
        else:
            # 取前 max_files 个文件
            all_files = all_files[:max_files]

    # 检查是否是Sleep-EDF数据集，如果是则按受试者分组
    is_sleep_edf = "sleep-edf" in data_dir.lower() or "edf" in data_dir.lower() or any("SC4" in os.path.basename(f) for f in all_files[:min(10, len(all_files))])
    if is_sleep_edf:
        print(f"检测到可能是Sleep-EDF数据集，使用按受试者分组的K折交叉验证")
        # 使用按受试者分组的k折交叉验证函数
        folds = split_files_by_subject_kfold(all_files, k=n_folds, data_path=data_dir, seed=42)
    else:
        # 使用普通的k折交叉验证函数
        print(f"使用普通K折交叉验证")
        folds = split_files_by_kfold(all_files, k=n_folds, shuffle=True, seed=42)

    return folds


def get_domain_files(source_dir, target_dir, max_files_source=None, max_files_target=None, add_random=False):
    """
    获取源域和目标域的数据文件列表
    """
    # 设置随机种子，确保文件选择的可复现性
    set_random_seed(42)
    
    # 使用工具类中的resolve_path函数处理路径
    source_dir = resolve_path(source_dir)
    target_dir = resolve_path(target_dir)
    
    # 获取源域文件
    source_files = glob(os.path.join(source_dir, "*.npz"))
    print(f"Found {len(source_files)} source domain files in {source_dir}")
    source_files.sort()
    
    # 获取目标域文件
    target_files = glob(os.path.join(target_dir, "*.npz"))
    print(f"Found {len(target_files)} target domain files in {target_dir}")
    target_files.sort()
    
    # 限制源域文件数量
    if max_files_source is not None and max_files_source > 0:
        if add_random:
            import random
            source_files = random.sample(source_files, min(max_files_source, len(source_files)))
        else:
            source_files = source_files[:max_files_source]
    
    # 限制目标域文件数量
    if max_files_target is not None and max_files_target > 0:
        if add_random:
            import random
            target_files = random.sample(target_files, min(max_files_target, len(target_files)))
        else:
            target_files = target_files[:max_files_target]
    
    return source_files, target_files
